#!/usr/bin/python3
# coding=utf-8
#
# Copyright (C) 2025. Huawei Technologies Co., Ltd. All rights reserved.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# ===============================================================================

import numpy as np
import os

def custom_maskfill(input_x, mask, value):
    """自定义MaskFill实现"""
    output = np.copy(input_x)
    output[mask] = value 
    return output


def gen_golden_data_maskfill():
    # 创建目录
    input_dir = "./input"
    output_dir = "./output"
    os.makedirs(input_dir, exist_ok=True)
    os.makedirs(output_dir, exist_ok=True)
    
    # 配置生成参数
    input_shape = [8, 2048]       # 输入张量形状
    input_dtype = np.float16      # 输入张量数据类型
    input_range = (-3, 3)         # 输入张量数值范围
    mask_true_ratio = 0.4         # 掩码中True的比例
    fill_value = -1.0             # 填充值
    
    input_x = np.random.uniform(low=input_range[0], high=input_range[1], size=input_shape).astype(input_dtype)
    mask = np.random.choice(a=[False, True], size=input_shape, p=[1 - mask_true_ratio, mask_true_ratio]).astype(np.bool_)
    fill_value_casted = np.array(fill_value, dtype=input_dtype)
    golden = custom_maskfill(input_x, mask, fill_value_casted)
    
    input_x.tofile(os.path.join(input_dir, "input_x.bin"))
    mask.tofile(os.path.join(input_dir, "mask.bin"))
    fill_value_casted.tofile(os.path.join(input_dir, "fill_value.bin"))
    golden.tofile(os.path.join(output_dir, "golden.bin"))

if __name__ == "__main__":
    gen_golden_data_maskfill()
